package mr

import (
	"log"
	"net"
	"net/http"
	"net/rpc"
	"os"
	"sync"
	"time"
)

type Master struct {
	// Your definitions here.
	nMap    int
	nReduce int
	mu      sync.Mutex

	// 1 -- Not given
	// 2 -- Given and not finished
	// 3 -- Finished
	inputList       map[string]int
	mapStartTime    map[string]time.Time
	givenMapTask    int
	finishedMapTask int
	inputIndex      map[string]int

	// 1 -- Not given
	// 2 -- Given and not finished
	// 3 -- Finished
	reduceList         map[int]int
	reduceStartTime    map[int]time.Time
	givenReduceTask    int
	finishedReduceTask int
}

// Your code here -- RPC handlers for the worker to call.
func (m *Master) GiveTask(args *AskTaskArgs, reply *AskTaskReply) error {
	m.mu.Lock()
	reply.NReduce = m.nReduce
	reply.NMap = m.nMap
	if m.givenMapTask < m.nMap {
		reply.IsMap = true
		reply.AllFinished = false
		for key, v := range m.inputList {
			if v == 1 {
				m.inputList[key] = 2
				m.givenMapTask += 1
				reply.InputFileName = key
				reply.TaskNumber = m.inputIndex[key]
				m.mapStartTime[key] = time.Now()
				defer m.mu.Unlock()
				return nil
			}
		}
	} else if m.finishedMapTask == m.nMap && m.givenReduceTask < m.nReduce {
		reply.IsMap = false
		reply.AllFinished = false
		for key, v := range m.reduceList {
			if v == 1 {
				m.reduceList[key] = 2
				m.givenReduceTask += 1
				reply.TaskNumber = key
				m.reduceStartTime[key] = time.Now()
				defer m.mu.Unlock()
				return nil
			}
		}
	}
	m.mu.Unlock()
	reply.AllGiven = true
	if m.finishedMapTask == m.nMap && m.finishedReduceTask == m.nReduce {
		reply.AllFinished = true
	}
	// return fmt.Errorf("all tasks are given, duplicate asking")
	return nil
}

func (m *Master) CrashCheck() {
	for {
		time.Sleep(5 * time.Second)
		currentTime := time.Now()
		m.mu.Lock()
		for k := range m.inputIndex {
			if m.inputList[k] == 2 && currentTime.Sub(m.mapStartTime[k]).Seconds() > 10 {
				// assume the work is dead
				// fmt.Printf("Map task %v crashed.\n", m.inputIndex[k])
				m.inputList[k] = 1
				m.givenMapTask -= 1
			}
		}
		for k, v := range m.reduceList {
			if v == 2 && currentTime.Sub(m.reduceStartTime[k]).Seconds() > 10 {
				// assume the work is dead
				// fmt.Printf("Reduce task %v crashed.\n", k)
				// err := os.Remove(fmt.Sprintf("mr-out-%v", k-1))
				// if err != nil {
				//   	log.Fatalf("cannot delete mr-out-%v", k-1)
				// }
				m.reduceList[k] = 1
				m.givenReduceTask -= 1
			}
		}
		m.mu.Unlock()
	}
}

func (m *Master) HandleFinishReport(args *FinishReportArgs, reply *FinishReportReply) error {
	m.mu.Lock()
	if args.IsFinishMap {
		for k, v := range m.inputIndex {
			if v == args.TaskNumber {
				if m.inputList[k] == 1 {
					// crashed, sadly
					defer m.mu.Unlock()
					return nil
				}
				m.finishedMapTask += 1
				m.inputList[k] = 3
				// fmt.Printf("Map task %v finished. Now finished %v / %v\n", k, m.finishedMapTask, m.nMap)
				defer m.mu.Unlock()
				return nil
			}
		}
	} else {
		if m.reduceList[args.TaskNumber] == 1 {
			// crashed, sadly
			defer m.mu.Unlock()
			return nil
		}
		m.finishedReduceTask += 1
		m.reduceList[args.TaskNumber] = 3
		// fmt.Printf("Reduce task %v finished. Now finished %v / %v\n", args.TaskNumber, m.finishedReduceTask, m.nReduce)
		defer m.mu.Unlock()
		return nil
	}
	m.mu.Unlock()
	return nil
}

// Example
// an example RPC handler.
//
// the RPC argument and reply types are defined in rpc.go.
//
func (m *Master) Example(args *ExampleArgs, reply *ExampleReply) error {
	reply.Y = args.X + 1
	return nil
}

//
// start a thread that listens for RPCs from worker.go
//
func (m *Master) server() {
	rpc.Register(m)
	rpc.HandleHTTP()
	//l, e := net.Listen("tcp", ":1234")
	sockname := masterSock()
	os.Remove(sockname)
	l, e := net.Listen("unix", sockname)
	if e != nil {
		log.Fatal("listen error:", e)
	}
	go http.Serve(l, nil)
}

// Done
// main/mrmaster.go calls Done() periodically to find out
// if the entire job has finished.
//
func (m *Master) Done() bool {
	ret := false

	// Your code here.
	if m.finishedMapTask == m.nMap && m.finishedReduceTask == m.nReduce {
		ret = true
	}

	return ret
}

// MakeMaster
// create a Master.
// main/mrmaster.go calls this function.
// nReduce is the number of reduce tasks to use.
//
func MakeMaster(files []string, nReduce int) *Master {
	m := Master{
		finishedMapTask:    0,
		finishedReduceTask: 0,
		nMap:               len(files),
		nReduce:            nReduce,
		givenMapTask:       0,
		givenReduceTask:    0,
		inputList:          make(map[string]int),
		reduceList:         make(map[int]int),
		inputIndex:         make(map[string]int),
		mapStartTime:       make(map[string]time.Time),
		reduceStartTime:    make(map[int]time.Time),
	}

	// Your code here.
	for i, fileName := range files {
		m.inputList[fileName] = 1
		m.inputIndex[fileName] = i
	}

	for i := 1; i <= nReduce; i++ {
		m.reduceList[i] = 1
	}

	go m.CrashCheck()

	m.server()
	return &m
}
